#include "stdafx.h"
#include "XepdtFT.h"
#include "FileUtil.h"
#include "log.h"
#include "CodeLib.h"

#define BLOCK_SIZE 10240

using namespace std;

CXepdtFT::CXepdtFT(Client* client_): m_client(client_)
, sendThread(this, &CXepdtFT::loopSend)
, recvThread(this, &CXepdtFT::loopRecv)
{
	m_ft = NULL;
	recvfile = NULL;
	m_bsSend = NULL;
	m_bsRecv = NULL;
	m_server = NULL;
	m_waitingSendLock = false;
	m_waitingRecvLock = false;
	SENDING_THREAD_RUNNING = false;
	RECEIVE_THREAD_RUNNING = false;

	m_ft = new SIProfileFT(m_client, this);

	HMODULE module = ::GetModuleHandle(NULL); 
	wchar_t pFileName[MAX_PATH]; 
	::GetModuleFileNameW(module, pFileName, MAX_PATH);
	std::wstring str = pFileName;

	m_testplanPath = str.substr(0, str.find_last_of(L"\\"));
	m_testplanPath += L"\\testplan\\";

	m_datafilePath = str.substr(0, str.find_last_of(L"\\"));
	m_datafilePath += L"\\testdata\\";
}

CXepdtFT::~CXepdtFT()
{
	if(SENDING_THREAD_RUNNING)
	{
		SENDING_THREAD_RUNNING = false;
		WaitForSingleObject(sendThread.handle_, INFINITE );
	}

	if(RECEIVE_THREAD_RUNNING)
	{
		RECEIVE_THREAD_RUNNING = false;
		WaitForSingleObject(recvThread.handle_, INFINITE);
	}

	delete m_server;
	delete m_ft;
	delete recvfile;
	m_client = NULL;
	m_ft = NULL;
	recvfile = NULL;
}

void* CXepdtFT::loopSend(void* )
{
	debug.log(CLog::Info, "thread to send datafile started");
	ConnectionError se = ConnNoError;
	char input[BLOCK_SIZE];
	bool isAnyBsSendingData = true;

	while(SENDING_THREAD_RUNNING)
	{
		//other thread is waiting for m_s5bRecvLock
		//Sleep to give a chance for it to get the lock
		
		if(m_waitingSendLock)
		{
			Sleep(5);	
		}

		m_s5bSendLock.lock();

		isAnyBsSendingData = false;

		FILELIST::iterator it = m_sfilelist.begin();

		for (; it != m_sfilelist.end();) 
		{
			FILEINFO info = (*it).second;
			std::fstream* sendfile = info.fileStream;

			if ( !sendfile->eof())
			{
				it++;			

				Bytestream* bs = info.bs;
				if (bs != NULL)
					bs->recv(1);

				if (bs != NULL && bs->isOpen()) 
				{
					sendfile->read(input, BLOCK_SIZE);
					std::string content(input, sendfile->gcount());

					bs->send(content);
					bs->recv(1);

					isAnyBsSendingData = true;
				}
					
			} 
			else 
			{
				Bytestream* bs = (*it).second.bs;
				if(bs != NULL)
				{
					bs->close();
					//The current element has been erase from m_sfilelist. See handleBytestreamClose
					//As a result the iterator is of no use, break from the current for loop
					break;
				}

				m_sfilelist.erase(it++);		
			}	
		}

		//If not sendint any data in the last loop, Sleep to avoid high CPU usage 
		if(!isAnyBsSendingData)
		{
			Sleep(5);
		}

		m_s5bSendLock.unlock();
	}

	debug.log(CLog::Info, "thread to send datafile ends");
	_endthreadex(0);

	return NULL;
}

void *CXepdtFT::loopRecv(void *)
{
	debug.log(CLog::Info, "thread to receive datafile start");
	while (RECEIVE_THREAD_RUNNING) 
	{
		//other thread is waiting for m_s5bRecvLock
		//Sleep to give a chance for it to get the lock
		//or there is no file to send, Sleep to avoid high CPU usage 
		if(m_waitingRecvLock)
		{
			Sleep(10);
		}

		m_s5bRecvLock.lock();
		if(m_rfilelist.size() == 0)
		{
			Sleep(10);
		}
		FILELIST::iterator it;

		for (it = m_rfilelist.begin(); it != m_rfilelist.end();) 
		{
			Bytestream* bs = (*it).second.bs;
			if(bs != NULL)
			  bs->recv(100);

			if((*it).second.transfered >= (*it).second.size)
			{
				(*it).second.fileStream->close();
				delete (*it).second.fileStream;
				m_rfilelist.erase(it++);
			}
			else
			{
				it++;
			}
		}
		m_s5bRecvLock.unlock();
	}
	debug.log(CLog::Info, "thread to receive datafile end");
	_endthreadex(0);

	return NULL;
}

bool CXepdtFT::isSend(Bytestream* bs)
{
	return bs->initiator().bare() == m_client->jid().bare();
}

bool CXepdtFT::ibbSendFile(const std::string& to, const std::string& file, const std::string& desc, std::string& sid)
{
	if(file.length() == 0 || sid.length() < 6)
	{
		return false;
	}
	//lock to ensure only one file can be sent at the same time
	m_ibbSend.lock();

	bool result = false;

	do
	{
		struct _stat f_stat;

		std::wstring unicodeName = CCodeLib::UTF8ToUnicode(file.c_str());
		if (_wstat(unicodeName.c_str(), &f_stat))
			break;

		int size = f_stat.st_size;

		std::fstream* sendfile = new std::fstream();

		std::string filename = FileUtil::getfilename(file);

		locale loc = locale::global(locale(""));
		sendfile->open(unicodeName.c_str(), std::ios_base::in | std::ios_base::binary);
		locale::global(loc);

		m_sid = m_ft->requestFT(JID(to), filename, size, EmptyString, desc, EmptyString, EmptyString, SIProfileFT::FTTypeIBB, JID(), sid);

		char input[BLOCK_SIZE];
		while ( !sendfile->eof() ) 
		{
			if(m_bsSend != NULL && m_bsSend->isOpen())
			{
				sendfile->read(input, BLOCK_SIZE);
				std::string content(input, sendfile->gcount());

				if (m_bsSend->send(content))
				{
				}
				m_bsSend->recv(1);
			}
		}
		delete sendfile;
		
		m_bsSend->removeBytestreamDataHandler();
		m_bsSend->close();

		result = true;
	}while(0);

	m_ibbSend.unlock();

	return result;
}

void CXepdtFT::startSock5Server(int port)
{
	SOCKS5BytestreamServer *pSock5 = new SOCKS5BytestreamServer( m_client->logInstance(), port );
	pSock5->listen();

	m_ft->registerSOCKS5BytestreamServer(pSock5);
	m_server = pSock5;

	std::string clientIP = m_client->connectionImpl()->localInterface();
	if(clientIP.length() != 0)
	{
		m_ft->addStreamHost( m_client->jid(), clientIP, port );
		debug.log(CLog::Info, "addStreamHost ip:%s, port:%d", clientIP.c_str(), port);
	}

	std::string serverIP = m_client->server();
	if(serverIP.length() != 0)
	{
		m_ft->addStreamHost( JID(serverIP + ":1090"), serverIP, 1090 );
		debug.log(CLog::Info, "addStreamHost ip:%s, port:%d", serverIP.c_str(), 1090);
	}
}

bool CXepdtFT::sock5SendFile(const std::string& to, const std::string& file, const std::string& desc, std::string& sid)
{
	if(file.length() == 0 || sid.length() < 6)
	{
		return false;
	}

	struct _stat f_stat;

	std::wstring unicodeName = CCodeLib::UTF8ToUnicode(file.c_str());
	if (_wstat(unicodeName.c_str(), &f_stat))
		return false;

	int size = f_stat.st_size;

	std::fstream* sendfile = new std::fstream();

	std::string filename = FileUtil::getfilename(file);

	locale loc = locale::global(locale(""));
	sendfile->open(unicodeName.c_str(), std::ios_base::in | std::ios_base::binary);
	locale::global(loc);

	string retSid = m_ft->requestFT(JID(to), filename, size, EmptyString, desc, EmptyString, EmptyString, SIProfileFT::FTTypeS5B, JID(), sid);

	m_waitingSendLock = true;
	m_s5bSendLock.lock();
	FILEINFO fInfo;
	fInfo.fileStream = sendfile;
	fInfo.size = size;
	fInfo.bs = NULL;
	fInfo.transfered = 0;
	m_sfilelist[retSid] = fInfo;
	m_s5bSendLock.unlock();
	m_waitingSendLock = false;

	if (!SENDING_THREAD_RUNNING)
	{
		SENDING_THREAD_RUNNING = true;
		sendThread.start();
	}

	return true;
}

void CXepdtFT::handleFTBytestream(Bytestream * bs)
{	
	if(bs == NULL)
	{
		debug.log(CLog::Error, "parameter bs is null for handleFTBytestream");
		return;
	}
	if(isSend(bs))
	{
		if(bs->type() == Bytestream::IBB)
		{
			m_bsSend = bs;
		}
		else if(bs->type() == Bytestream::S5B)
		{
			m_waitingSendLock = true;
			m_s5bSendLock.lock();
			FILELIST::iterator iter = m_sfilelist.find(bs->sid());
			if(iter != m_sfilelist.end())
			{
				FILEINFO& fInfo = (*iter).second;
				fInfo.bs = bs;
			}
			//m_bssendList.push_back(bs);
			m_s5bSendLock.unlock();

			m_waitingSendLock = false;
		}
	}
	else 
	{
		if(bs->type() == Bytestream::IBB)
		{
			m_bsRecv = bs;
		}
		else if(bs->type() == Bytestream::S5B)
		{
			m_waitingRecvLock = true;
			m_s5bRecvLock.lock();
			FILELIST::iterator iter = m_rfilelist.find(bs->sid());
			if(iter != m_rfilelist.end())
			{
				FILEINFO& fInfo = (*iter).second;
				fInfo.bs = bs;
			}
			m_s5bRecvLock.unlock();
			m_waitingRecvLock = false;

			if(!RECEIVE_THREAD_RUNNING)
			{
				RECEIVE_THREAD_RUNNING = true;
				recvThread.start();
			}
		}
	}

	bs->registerBytestreamDataHandler(this);

	if (bs->connect())
	{
		debug.log(CLog::Info, "ByteStream connection success");
	}
	else
	{
		debug.log(CLog::Error, "ByteStream connection failed");
	}
}


void CXepdtFT::handleFTRequest( const JID& from, const JID& to, const std::string& sid,
							   const std::string& name, long size, const std::string& hash,
							   const std::string& date, const std::string& mimetype,
							   const std::string& desc, int stypes )
{
	m_ft->acceptFT(from, sid, (SIProfileFT::StreamType)stypes); //SIProfileFT::FTTypeIBB);

	if(stypes == SIProfileFT::FTTypeIBB)
	{
		m_sid = sid;

		recvfile = new std::fstream();

		locale loc = locale::global(locale(""));
		wstring filename;
		if(m_currentTestPlanName.length() == 0)
		{
			filename = m_testplanPath + CCodeLib::UTF8ToUnicode(name.c_str());
		}
		else
		{
			filename = m_currentTestPlanName;
		}
		recvfile->open(filename.c_str(), std::ios_base::out | std::ios_base::binary);
		locale::global(loc);

		debug.log(CLog::Info, "IBB start to receive file:%s", name.c_str());
	}

	if(stypes == SIProfileFT::FTTypeS5B)
	{
		std::fstream* recvstream = new std::fstream();

		wstring filename;
		if(m_currentDataFileName.length() == 0)
		{
			filename = m_datafilePath + CCodeLib::UTF8ToUnicode(name.c_str());
		}
		else
		{
			filename = m_currentDataFileName;
		}
		recvstream->open(filename.c_str(), std::ios_base::out | std::ios_base::binary);

		m_waitingRecvLock = true;
		m_s5bRecvLock.lock();

		FILEINFO fInfo;
		fInfo.fileStream = recvstream;
		fInfo.filename = CCodeLib::UnicodeToUTF8(filename.c_str());
		fInfo.size = size;
		fInfo.transfered = 0;
		fInfo.bs = NULL;
		m_rfilelist[sid] = fInfo;

		m_s5bRecvLock.unlock();
		m_waitingRecvLock = false;

		debug.log(CLog::Info, "Sock5 start to receive file:%s", name.c_str());
	}
}

void CXepdtFT::handleBytestreamOpen(Bytestream * bs)
{
}

void CXepdtFT::handleBytestreamClose(Bytestream * bs)
{
	if(bs == NULL)
		return;

	if(bs->type() == Bytestream::S5B)
	{
		if (isSend(bs))
		{
			m_waitingSendLock = true;

			m_s5bSendLock.lock();
			FILELIST::iterator s_iter = m_sfilelist.find(bs->sid());

			if (s_iter != m_sfilelist.end()) 
			{
				(*s_iter).second.fileStream->close();
				delete (*s_iter).second.fileStream;
				m_sfilelist.erase(s_iter);
			}
			m_s5bSendLock.unlock();

			m_waitingSendLock = false;

			bs->removeBytestreamDataHandler();
			bs->close();
		} 
		else 
		{
			/////////////////////////////////////////////
			//When we receive stream close message, the data may still be in the progress of transfering
			//So we do the clean up in handleBytestreamData instead of here
		}
	}

	else if(bs->type() == Bytestream::IBB)
	{
		bs->removeBytestreamDataHandler();
		bs->close();

		m_sid.clear();
		recvfile->close();
		recvfile = NULL;
	}
}

void CXepdtFT::handleBytestreamError(Bytestream * bs, const IQ & stanza)
{
	handleBytestreamClose(bs);
}

void CXepdtFT::handleBytestreamData(Bytestream * bs,
									const std::string & data)
{
	if(bs->type() == Bytestream::S5B)
	{
		FILELIST::iterator iter = m_rfilelist.find(bs->sid());
		(*iter).second.fileStream->write(data.c_str(), data.length());
		(*iter).second.transfered += data.length();
		if((*iter).second.transfered == (*iter).second.size)
		{
			bs->removeBytestreamDataHandler();
			bs->close();
			debug.log(CLog::Info, "Receive file finished");
		}
		else if((*iter).second.transfered > (*iter).second.size)
		{
			debug.log(CLog::Error, "Received size is larger than file size");
			bs->removeBytestreamDataHandler();
			bs->close();
		}

		string client = (*iter).second.bs->initiator().bare();
		mdt_file file;
		file.filename = (*iter).second.filename;
		file.size = (*iter).second.size;
		file.filetype = eFT_DataFile;

		float percentage = (float)(*iter).second.transfered/(*iter).second.size;
		mdt_transfer transfer;
		transfer.per = (int)(percentage * 100);
		transfer.status = eTS_Transfer;
		if (m_fDoTransfer)
		{
			m_fDoTransfer(client, file, transfer);
		}
	}
	else if(bs->type() == Bytestream::IBB)
	{
		//only one testplan can be send at the same time. So it is very safe to write the data to the 
		//file stream directly
		recvfile->write(data.c_str(), data.length());
	}
}

void CXepdtFT::handleFTRequestError(const IQ & iq, const std::string & sid)
{
}


void CXepdtFT::setCurrentTestPlanName(const std::string& desc)
{
	m_currentTestPlanName = CCodeLib::UTF8ToUnicode(desc.c_str());
}

void CXepdtFT::setCurrentDataFileName(const std::string& desc)
{
	m_currentDataFileName = CCodeLib::UTF8ToUnicode(desc.c_str());
}

void CXepdtFT::SetTransferHandler(const TransferHandler& fDoTransfer)
{
	m_fDoTransfer = fDoTransfer;
}